import torch
import torch.nn as nn
from torch.nn import functional as F
from efficientnet import efficientnet_b4

class ModalitySpecificNet(nn.Module):
    def __init__(self, in_channels=None, num_classes=None):
        super().__init__()
        self.effnet = efficientnet_b4(in_channels=in_channels, num_classes=num_classes)

    def forward(self, img):
        enhanced_img = img
        features = self.effnet(enhanced_img)
        return features

class MultiModalEfficientNet(nn.Module):
    def __init__(self, in_channels=None, num_classes=None):
        super().__init__()
        # 三个模态的独立网络
        self.axi_net = ModalitySpecificNet(in_channels=in_channels, num_classes=num_classes)
        self.cor_net = ModalitySpecificNet(in_channels=in_channels, num_classes=num_classes)
        self.sag_net = ModalitySpecificNet(in_channels=in_channels, num_classes=num_classes)
        
        # 特征融合层 - 使用通道注意力机制（更新为448通道）
        self.channel_attention = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(448 * 3, 448 // 16, 1, bias=False),  # 1344 -> 28
            nn.ReLU(inplace=True),
            nn.Conv2d(448 // 16, 448 * 3, 1, bias=False),   # 28 -> 1344
            nn.Sigmoid()
        )
        
        # 空间注意力机制
        self.spatial_attention = nn.Sequential(
            nn.Conv2d(3, 1, kernel_size=7, padding=3, bias=False),
            nn.Sigmoid()
        )
        
        # 最终分类层（根据448通道调整）
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(448, 256),  # 增加中间层维度以适应更多特征
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(256, 128),
            nn.ReLU(inplace=True),
            nn.Dropout(0.3),
            nn.Linear(128, num_classes if num_classes else 1)
        )

    def forward(self, multimodal_img):
        axi = multimodal_img[:, 0, :, :, :]  # 提取 axial 模态
        cor = multimodal_img[:, 1, :, :, :]  # 提取 coronal 模态
        sag = multimodal_img[:, 2, :, :, :]  # 提取 sagittal 模态
        
        # 提取各模态特征 [4, 448, 4, 4]
        axi_features = self.axi_net(axi)
        cor_features = self.cor_net(cor)
        sag_features = self.sag_net(sag)
        
        # 通道维度拼接 [4, 1344, 4, 4]
        concatenated = torch.cat([axi_features, cor_features, sag_features], dim=1)
        
        # 通道注意力
        channel_att = self.channel_attention(concatenated)
        attended_features = concatenated * channel_att
        
        # 分离回各模态特征
        axi_att = attended_features[:, :448, :, :]
        cor_att = attended_features[:, 448:896, :, :]
        sag_att = attended_features[:, 896:, :, :]
        
        # 空间注意力（基于各模态特征的重要性）
        spatial_weights = torch.stack([
            axi_att.mean(dim=1, keepdim=True),
            cor_att.mean(dim=1, keepdim=True), 
            sag_att.mean(dim=1, keepdim=True)
        ], dim=1)  # [4, 3, 1, 4, 4]
        
        spatial_att = self.spatial_attention(spatial_weights.squeeze(2))  # [4, 1, 4, 4]
        
        # 加权融合
        fused_features = (axi_att + cor_att + sag_att) * spatial_att  # [4, 448, 4, 4]
        
        # 分类
        output = self.classifier(fused_features)
        
        return output
        
